/****************************************************************
* @file   : detection.c
* @author : NingJian (freegodly@gmail.com)
* @date   : 06/09/21 21:06:26
* @brief  : 
****************************************************************/

#include <detection.h>

CMat *detection_forwad(CMat *p,
                        CMat *anchors,
                        int nA,
                        int nC,
                        int stride,
                        int training)
{
    /***************************************************************************
    bs, nG = p.shape[0], p.shape[-1]
    p = p.view(bs, self.__nA, 5 + self.__nC, nG, nG).permute(0, 3, 4, 1, 2)
    **************************************************************************/
    int bs = p->dim_size[0];
    int nG = p->dim_size[p->dim - 1];
    int new_p_info[5] = {bs, nA, 5 + nC, nG, nG};
    CMat_view(p, 5, new_p_info);
    //CMat_print(p);
    int permute_info[5] = {0, 3, 4, 1, 2};
    CMat *new_p = CMat_permute(p, permute_info);
    //CMat_print(new_p);

    CMat *p_clone = CMat_clone(new_p);
    //CMat_print(p_clone);

    /***************************************************************************
    batch_size, output_size = p.shape[:2]

    device = p.device
    stride = self.__stride
    anchors = (1.0 * self.__anchors).to(device)

    conv_raw_dxdy = p[ :, :, :, :, 0:2]
    conv_raw_dwdh = p[ :, :, :, :, 2:4]
    conv_raw_conf = p[ :, :, :, :, 4:5]
    conv_raw_prob = p[ :, :, :, :, 5:]
    **************************************************************************/
    int batch_size = p_clone->dim_size[0];
    int output_size = p_clone->dim_size[1];
    //printf("batch_size = %d  output_size = %d\r\n", batch_size, output_size);

    int pick_info[5][2] = {{0, p_clone->dim_size[0]},
                           {0, p_clone->dim_size[1]},
                           {0, p_clone->dim_size[2]},
                           {0, p_clone->dim_size[3]},
                           {0, 2}};
    CMat *conv_raw_dxdy = CMat_pick(p_clone, pick_info);

    pick_info[4][0] = 2;
    pick_info[4][1] = 4;
    CMat *conv_raw_dwdh = CMat_pick(p_clone, pick_info);

    pick_info[4][0] = 4;
    pick_info[4][1] = 5;
    CMat *conv_raw_conf = CMat_pick(p_clone, pick_info);

    pick_info[4][0] = 5;
    pick_info[4][1] = p_clone->dim_size[4];
    CMat *conv_raw_prob = CMat_pick(p_clone, pick_info);

    CMat_delete(p_clone);

    /***************************************************************************
    y = torch.arange(0, output_size).unsqueeze(1).repeat(1, output_size)
    x = torch.arange(0, output_size).unsqueeze(0).repeat(output_size, 1)
    grid_xy = torch.stack([x, y], dim = -1)
    grid_xy = grid_xy.unsqueeze(0).unsqueeze(3).repeat(batch_size, 1, 1, 3, 1).double()
    **************************************************************************/

    CMat *x = CMat_arange_unsqueeze_repeat(output_size, 0);
    CMat *y = CMat_arange_unsqueeze_repeat(output_size, 1);

    CMat *grid_xy = CMat_stack(x, y, -1);

    CMat_unsqueeze(grid_xy, 0);
    CMat_unsqueeze(grid_xy, 3);
    int repeat_info[5] = {batch_size, 1, 1, 3, 1};
    CMat *new_grid_xy = CMat_repeat(grid_xy, repeat_info);

    //CMat_print(new_grid_xy);
    CMat_delete(x);
    CMat_delete(y);
    CMat_delete(grid_xy);
    grid_xy = new_grid_xy;

    /***************************************************************************
    pred_xy = (torch.sigmoid(conv_raw_dxdy) + grid_xy) * stride
    pred_wh = (torch.exp(conv_raw_dwdh) * anchors) * stride
    pred_xywh = torch.cat([pred_xy, pred_wh], dim = -1)
    pred_conf = torch.sigmoid(conv_raw_conf)
    pred_prob = torch.sigmoid(conv_raw_prob)
    pred_bbox = torch.cat([pred_xywh, pred_conf, pred_prob], dim = -1)
    **************************************************************************/

    CMat_sigmoid(conv_raw_dxdy);
    CMat_add_mat(conv_raw_dxdy, grid_xy);
    CMat_mult_value(conv_raw_dxdy, stride);
    CMat *pred_xy = conv_raw_dxdy;
    conv_raw_dxdy = NULL;

    CMat_exp(conv_raw_dwdh);
    CMat_mult_mat(conv_raw_dwdh, anchors);
    CMat_mult_value(conv_raw_dwdh, stride);
    CMat *pred_wh = conv_raw_dwdh;
    conv_raw_dwdh = NULL;

    CMat *mats[2] = {pred_xy, pred_wh};
    CMat *pred_xywh = CMat_cat(mats, 2, -1);

    CMat_sigmoid(conv_raw_conf);
    CMat *pred_conf = conv_raw_conf;
    conv_raw_conf = NULL;

    CMat_sigmoid(conv_raw_prob);
    CMat *pred_prob = conv_raw_prob;
    conv_raw_prob = NULL;

    CMat *mats_bbox[3] = {pred_xywh, pred_conf, pred_prob};
    CMat *pred_bbox = CMat_cat(mats_bbox, 3, -1);

    CMat_delete(grid_xy);
    CMat_delete(pred_xy);
    CMat_delete(pred_wh);
    CMat_delete(pred_xywh);
    CMat_delete(pred_conf);
    CMat_delete(pred_prob);

    /***************************************************************************
    return pred_bbox.view(-1, 5 + self.__nC) if not self.training else pred_bbox
    **************************************************************************/

    int view_dim_size[2] = {pred_bbox->data_len / (5 + nC), 5 + nC};
    if (!training)
    {
        CMat_view(pred_bbox, 2, view_dim_size);
    }
    //CMat_print(pred_bbox);
//    static CMat *result[2];
//    result[0] = new_p;
//    result[1] = pred_bbox;

    CMat_delete(new_p);
    return pred_bbox;
}